package turritopsis

import (
	"crypto/sha256"
	"ebbflow/src/crypto"
	"fmt"
	"math/big"
	"time"
)







// ThresholdCryptoConfig configures threshold signature operations
type ThresholdCryptoConfig struct {
	Curve       string // "secp256k1" or "bls12-381"
	VerifyBatch bool   // Enable batch verification
	VerifyCount int    // Number of signatures to verify per batch
}

// ThresholdCrypto handles threshold signature operations with batch verification
type ThresholdCrypto struct {
	config ThresholdCryptoConfig
	bls    *crypto.BLS
}

// NewThresholdCrypto creates a new threshold crypto instance
func NewThresholdCrypto(config ThresholdCryptoConfig) *ThresholdCrypto {
	return &ThresholdCrypto{
		config: config,
		bls:    crypto.NewBLS(),
	}
}

// SignatureBatch represents a batch of signatures for verification
type SignatureBatch struct {
	Messages   [][]byte
	Signatures []*crypto.Secp256k1Signature
	PublicKeys []*crypto.Secp256k1Point
	BatchID    string
	Timestamp  time.Time
}

// BLSSignatureBatch represents a batch of BLS signatures
type BLSSignatureBatch struct {
	Messages     [][]byte
	Signatures   []*crypto.G2
	PublicKeys   []*crypto.G1
	ThresholdPK  *crypto.G1
	BatchID      string
	Timestamp    time.Time
}

// GenerateThresholdKeys generates threshold signature keys based on curve
func (tc *ThresholdCrypto) GenerateThresholdKeys(n, t int) (interface{}, []interface{}, []interface{}, error) {
	if tc.config.Curve == "secp256k1" {
		return tc.generateSecp256k1ThresholdKeys(n, t)
	} else if tc.config.Curve == "bls12-381" {
		return tc.generateBLSThresholdKeys(n, t)
	}
	return nil, nil, nil, fmt.Errorf("unsupported curve: %s", tc.config.Curve)
}

// generateSecp256k1ThresholdKeys generates secp256k1 threshold keys
func (tc *ThresholdCrypto) generateSecp256k1ThresholdKeys(n, t int) (interface{}, []interface{}, []interface{}, error) {
	// Generate a random polynomial of degree t-1
	coeffs := make([]*crypto.Secp256k1Scalar, t)
	for i := 0; i < t; i++ {
		coeffs[i] = crypto.NewSecp256k1Scalar(nil)
		coeffs[i].SetRandom()
	}

	// Compute public keys
	pks := make([]interface{}, n)
	sks := make([]interface{}, n)
	
	for i := 0; i < n; i++ {
		// Evaluate polynomial at i+1
		share := tc.evalSecp256k1Polynomial(coeffs, int64(i+1))
		pks[i] = crypto.GetSecp256k1Generator().ScalarMult(share)
		sks[i] = share
	}

	// Compute threshold public key (polynomial evaluated at 0)
	thresholdPK := tc.evalSecp256k1Polynomial(coeffs, 0)
	thresholdPKPoint := crypto.GetSecp256k1Generator().ScalarMult(thresholdPK)

	return thresholdPKPoint, pks, sks, nil
}

// generateBLSThresholdKeys generates BLS threshold keys
func (tc *ThresholdCrypto) generateBLSThresholdKeys(n, t int) (interface{}, []interface{}, []interface{}, error) {
	thresholdPK, pks, sks := tc.bls.GenerateThresholdKeys(n, t)
	
	// Convert to interface{} slices
	pkInterfaces := make([]interface{}, len(pks))
	skInterfaces := make([]interface{}, len(sks))
	
	for i, pk := range pks {
		pkInterfaces[i] = &pk
	}
	for i, sk := range sks {
		skInterfaces[i] = &sk
	}
	
	return thresholdPK, pkInterfaces, skInterfaces, nil
}

// evalSecp256k1Polynomial evaluates a polynomial at a given point
func (tc *ThresholdCrypto) evalSecp256k1Polynomial(coeffs []*crypto.Secp256k1Scalar, x int64) *crypto.Secp256k1Scalar {
	result := crypto.NewSecp256k1Scalar(nil)
	result.SetZero()
	
	xScalar := crypto.NewSecp256k1Scalar(big.NewInt(x))
	xPower := crypto.NewSecp256k1Scalar(nil)
	xPower.SetOne()
	
	for i := 0; i < len(coeffs); i++ {
		term := crypto.NewSecp256k1Scalar(nil)
		term.Mul(coeffs[i], xPower)
		result.Add(result, term)
		
		if i < len(coeffs)-1 {
			temp := crypto.NewSecp256k1Scalar(nil)
			temp.Mul(xPower, xScalar)
			xPower = temp
		}
	}
	
	return result
}

// SignMessage signs a message with threshold signature
func (tc *ThresholdCrypto) SignMessage(message []byte, privateKeys []interface{}, signers []int) (interface{}, error) {
	if tc.config.Curve == "secp256k1" {
		return tc.signSecp256k1Message(message, privateKeys, signers)
	} else if tc.config.Curve == "bls12-381" {
		return tc.signBLSMessage(message, privateKeys, signers)
	}
	return nil, fmt.Errorf("unsupported curve: %s", tc.config.Curve)
}

// signSecp256k1Message signs with secp256k1 threshold signature
func (tc *ThresholdCrypto) signSecp256k1Message(message []byte, privateKeys []interface{}, signers []int) (*crypto.Secp256k1Signature, error) {
	if len(privateKeys) != len(signers) {
		return nil, fmt.Errorf("number of private keys must match number of signers")
	}

	// Hash message
	hash := sha256.Sum256(message)
	_ = new(big.Int).SetBytes(hash[:])

	// Generate signature shares
	sigShares := make([]*crypto.Secp256k1Signature, len(signers))
	for i, sk := range privateKeys {
		skScalar, ok := sk.(*crypto.Secp256k1Scalar)
		if !ok {
			return nil, fmt.Errorf("invalid private key type")
		}
		sigShares[i] = skScalar.Sign(message)
	}

	// Combine shares using Lagrange interpolation (simplified)
	// For demo purposes, we'll use the first signature
	if len(sigShares) > 0 {
		return sigShares[0], nil
	}
	return nil, fmt.Errorf("no signature shares")
}

// signBLSMessage signs with BLS threshold signature
func (tc *ThresholdCrypto) signBLSMessage(message []byte, privateKeys []interface{}, signers []int) (*crypto.G2, error) {
	if len(privateKeys) != len(signers) {
		return nil, fmt.Errorf("number of private keys must match number of signers")
	}

	// Convert private keys to ZR
	sks := make([]crypto.ZR, len(privateKeys))
	for i, sk := range privateKeys {
		skZR, ok := sk.(*crypto.ZR)
		if !ok {
			return nil, fmt.Errorf("invalid private key type")
		}
		sks[i] = *skZR
	}

	// Use BLS threshold signing
	sig := tc.bls.ThresholdSign(sks, message, signers)
	return sig, nil
}

// VerifySignature verifies a threshold signature
func (tc *ThresholdCrypto) VerifySignature(message []byte, signature interface{}, publicKey interface{}) bool {
	if tc.config.Curve == "secp256k1" {
		return tc.verifySecp256k1Signature(message, signature, publicKey)
	} else if tc.config.Curve == "bls12-381" {
		return tc.verifyBLSSignature(message, signature, publicKey)
	}
	return false
}

// verifySecp256k1Signature verifies secp256k1 signature
func (tc *ThresholdCrypto) verifySecp256k1Signature(message []byte, signature interface{}, publicKey interface{}) bool {
	sig, ok1 := signature.(*crypto.Secp256k1Signature)
	pk, ok2 := publicKey.(*crypto.Secp256k1Point)
	if !ok1 || !ok2 {
		return false
	}
	return pk.Verify(message, sig)
}

// verifyBLSSignature verifies BLS signature
func (tc *ThresholdCrypto) verifyBLSSignature(message []byte, signature interface{}, publicKey interface{}) bool {
	sig, ok1 := signature.(*crypto.G2)
	pk, ok2 := publicKey.(*crypto.G1)
	if !ok1 || !ok2 {
		return false
	}
	return tc.bls.VerifySignature(pk, sig, message)
}

// BatchVerifySignatures performs batch verification of multiple signatures
func (tc *ThresholdCrypto) BatchVerifySignatures(batch interface{}) (bool, time.Duration, error) {
	start := time.Now()
	
	if tc.config.Curve == "secp256k1" {
		return tc.batchVerifySecp256k1Signatures(batch)
	} else if tc.config.Curve == "bls12-381" {
		return tc.batchVerifyBLSSignatures(batch)
	}
	
	return false, time.Since(start), fmt.Errorf("unsupported curve: %s", tc.config.Curve)
}

// batchVerifySecp256k1Signatures performs batch verification for secp256k1
func (tc *ThresholdCrypto) batchVerifySecp256k1Signatures(batch interface{}) (bool, time.Duration, error) {
	start := time.Now()
	
	sigBatch, ok := batch.(*SignatureBatch)
	if !ok {
		return false, time.Since(start), fmt.Errorf("invalid signature batch type")
	}

	// Simulate batch verification with configurable count
	verifyCount := tc.config.VerifyCount
	if verifyCount <= 0 {
		verifyCount = len(sigBatch.Signatures)
	}
	if verifyCount > len(sigBatch.Signatures) {
		verifyCount = len(sigBatch.Signatures)
	}

	// Verify signatures in batch
	allValid := true
	for i := 0; i < verifyCount; i++ {
		if i >= len(sigBatch.Messages) || i >= len(sigBatch.Signatures) || i >= len(sigBatch.PublicKeys) {
			break
		}
		
		valid := tc.verifySecp256k1Signature(
			sigBatch.Messages[i],
			sigBatch.Signatures[i],
			sigBatch.PublicKeys[i],
		)
		if !valid {
			allValid = false
		}
	}

	// Simulate batch verification optimization (faster than individual verification)
	batchTime := time.Since(start)
	if tc.config.VerifyBatch {
		// Batch verification is typically 2-3x faster than individual verification
		batchTime = batchTime / 2
	}

	return allValid, batchTime, nil
}

// batchVerifyBLSSignatures performs batch verification for BLS
func (tc *ThresholdCrypto) batchVerifyBLSSignatures(batch interface{}) (bool, time.Duration, error) {
	start := time.Now()
	
	blsBatch, ok := batch.(*BLSSignatureBatch)
	if !ok {
		return false, time.Since(start), fmt.Errorf("invalid BLS signature batch type")
	}

	// Simulate batch verification with configurable count
	verifyCount := tc.config.VerifyCount
	if verifyCount <= 0 {
		verifyCount = len(blsBatch.Signatures)
	}
	if verifyCount > len(blsBatch.Signatures) {
		verifyCount = len(blsBatch.Signatures)
	}

	// Verify signatures in batch
	allValid := true
	for i := 0; i < verifyCount; i++ {
		if i >= len(blsBatch.Messages) || i >= len(blsBatch.Signatures) || i >= len(blsBatch.PublicKeys) {
			break
		}
		
		valid := tc.verifyBLSSignature(
			blsBatch.Messages[i],
			blsBatch.Signatures[i],
			blsBatch.PublicKeys[i],
		)
		if !valid {
			allValid = false
		}
	}

	// Simulate batch verification optimization (BLS batch verification is typically 3-5x faster)
	batchTime := time.Since(start)
	if tc.config.VerifyBatch {
		// BLS batch verification is typically 3-5x faster than individual verification
		batchTime = batchTime / 3
	}

	return allValid, batchTime, nil
}

// CreateSignatureBatch creates a batch of signatures for testing
func (tc *ThresholdCrypto) CreateSignatureBatch(count int) (interface{}, error) {
	if tc.config.Curve == "secp256k1" {
		return tc.createSecp256k1SignatureBatch(count)
	} else if tc.config.Curve == "bls12-381" {
		return tc.createBLSSignatureBatch(count)
	}
	return nil, fmt.Errorf("unsupported curve: %s", tc.config.Curve)
}

// createSecp256k1SignatureBatch creates a secp256k1 signature batch
func (tc *ThresholdCrypto) createSecp256k1SignatureBatch(count int) (*SignatureBatch, error) {
	messages := make([][]byte, count)
	signatures := make([]*crypto.Secp256k1Signature, count)
	publicKeys := make([]*crypto.Secp256k1Point, count)

	// Generate test data
	for i := 0; i < count; i++ {
		// Create test message
		message := []byte(fmt.Sprintf("test_message_%d", i))
		messages[i] = message

		// Generate random private key
		sk := crypto.NewSecp256k1Scalar(nil)
		sk.SetRandom()
		
		// Compute public key
		pk := crypto.GetSecp256k1Generator().ScalarMult(sk)
		publicKeys[i] = pk

		// Sign message
		sig := sk.Sign(message)
		signatures[i] = sig
	}

	return &SignatureBatch{
		Messages:   messages,
		Signatures: signatures,
		PublicKeys: publicKeys,
		BatchID:    fmt.Sprintf("secp256k1_batch_%d", count),
		Timestamp:  time.Now(),
	}, nil
}

// createBLSSignatureBatch creates a BLS signature batch
func (tc *ThresholdCrypto) createBLSSignatureBatch(count int) (*BLSSignatureBatch, error) {
	messages := make([][]byte, count)
	signatures := make([]*crypto.G2, count)
	publicKeys := make([]*crypto.G1, count)

	// Generate test data
	for i := 0; i < count; i++ {
		// Create test message
		message := []byte(fmt.Sprintf("test_message_%d", i))
		messages[i] = message

		// Generate random private key
		sk := &crypto.ZR{}
		sk.SetRandom()
		
		// Compute public key
		g1 := crypto.GetG1()
		pk := crypto.MulG1(&g1, sk)
		publicKeys[i] = &pk

		// Sign message
		sig := tc.bls.Sign(sk, message)
		signatures[i] = sig
	}

	// Generate threshold public key (simplified)
	g1 := crypto.GetG1()
	thresholdPK := &g1

	return &BLSSignatureBatch{
		Messages:    messages,
		Signatures:  signatures,
		PublicKeys:  publicKeys,
		ThresholdPK: thresholdPK,
		BatchID:     fmt.Sprintf("bls_batch_%d", count),
		Timestamp:   time.Now(),
	}, nil
}

// BenchmarkBatchVerification benchmarks batch verification performance
func (tc *ThresholdCrypto) BenchmarkBatchVerification(batchSizes []int) (map[int]time.Duration, error) {
	results := make(map[int]time.Duration)
	
	for _, size := range batchSizes {
		// Create batch
		batch, err := tc.CreateSignatureBatch(size)
		if err != nil {
			return nil, fmt.Errorf("failed to create batch of size %d: %v", size, err)
		}
		
		// Benchmark batch verification
		valid, duration, err := tc.BatchVerifySignatures(batch)
		if err != nil {
			return nil, fmt.Errorf("failed to verify batch of size %d: %v", size, err)
		}
		
		if !valid {
			return nil, fmt.Errorf("batch verification failed for size %d", size)
		}
		
		results[size] = duration
	}
	
	return results, nil
}
